#include "sweSolver.hpp"
#include "sweFunctions.hpp"

namespace hms
{

/** @brief Calculates the mass and momentum sources and subtracts them from the
 * storage field */
template<typename MeshType>
void SWESolver<MeshType>::calculateSourcesImp(){
	/* mass sources */
	// auto addSourcesBlockwise = [&]( Index colBegDomain, Index colBeg, Index nCols ){
	auto addSourcesBlockwise = [&]( Index colBeg, Index nCols ){

		auto qBlock = q().values().template topRows<1>().middleCols(colBeg, nCols);

		auto visitor = [&](const auto& sourceField){
			qBlock +=
				this->time().step() *
				// sourceField.values().middleCols(colBegDomain, nCols);
				sourceField.values().middleCols(colBeg, nCols);
		};

		this->visitUsedSourceFields( visitor );

		/* prevent negative water depths (may arise from infiltration) */
		qBlock = qBlock.max(0);
	};

	if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
		static constexpr Index ng {1};
		Index
			nx = this->mesh().nCellsX(),
			ny = this->mesh().nCellsY(),
			cellsPerRow { 2*ng + nx },
			domainCellsBeg { ng*(cellsPerRow+1) };
		
		OMP_PRAGMA( parallel for )
		for ( int j = 0; j < ny; ++j ){
			// addSourcesBlockwise( nx*j, domainCellsBeg + j*cellsPerRow, nx );
			addSourcesBlockwise( domainCellsBeg + j*cellsPerRow, nx );
		}
	} else {
		OMP_PRAGMA( parallel )
		{

		auto [colBeg, nCols] = ompSegment( this->mesh().nDomainCells() );

		addSourcesBlockwise(colBeg, nCols);

		/** @todo momentum sources */
		/* if fluxes are stored separately,
		 * momentum sources need to be written to fluxes,
		 * because preconditioning is applied to them */
		}
	}



	/* Friction source term with preconditioning */
	/* See Simons (2019), Eq. 3.109. The idea is to recalculate the storage
	 * vector from the state variables, so that the difference to the current 
	 * storage vector is equal to the sum of fluxes and sources. These can then
	 * be modified by the preconditioning process and added back onto the
	 * storage vector calculated from the state variables.
	 * Before modification by the preconditioner, the actual friction source
	 * term is added to the sum of fluxes and sources. */
	if ( this->settings().useFriction ){
		scalar dt { this->time().step() };

		auto applyFrictionBlockwise = [&](
			const auto& n, Index colBeg, Index nCols
		){
			auto qBlock = q().values()
				.template bottomRows<2>().middleCols(colBeg, nCols);
			auto dBlock = Array1Xs::Map( &d().values()(0, colBeg), 1, nCols );
			auto vBlock = Array2Xs::Map( &v().values()(0, colBeg), 2, nCols);
			// auto nBlock =   n.values().middleCols(colBeg, nCols);

			using nType =
				std::remove_cv_t<std::remove_reference_t<decltype( n.values() )>>;
			auto nBlock = [&](){
				if constexpr ( std::is_same_v<nType, ArrayXXs> )
					return Eigen::Map<const Array1Xs>(
						&n.values()(0, colBeg), nCols
					);
				else
					return n.values().middleCols(colBeg, nCols);
			}();

			Array1Xs vNorm = vBlock.matrix().colwise().norm().array();

			/* this factor is g / C^2, with Chezy coeff. C = d^(1/6) / n */
			Array1Xs chezyFac = phys::g * nBlock.pow(2) / dBlock.pow(1./3);

			Array1Xs fricFac = dt * (chezyFac * vNorm);
			chezyFac *= dt / dBlock;
			
			auto vNormRepl { vNorm.template replicate<2,1>() };
			auto fricRepl  { fricFac.template replicate<2,1>() };
			auto chezyRepl { chezyFac.template replicate<2,1>() };

			#ifdef SOLVER_STORE_FLUXES

			auto momentumFluxBlock = this->fluxes().values()
				.template bottomRows<2>().middleCols(colBeg, nCols);

			qBlock += ( dBlock < minWaterDepth || vNorm == 0 )
				.template replicate<2,1>().select(
					0,
					( momentumFluxBlock - fricRepl * vBlock ) /
					( 1 + chezyRepl * ( vNormRepl + vBlock.pow(2) / vNormRepl ) )
				);

			#else

			/* by subtracting the old storage values from the current ones, we
			 * get all the fluxes and sources to use for point-implicit 
			 * treatment. However, first, the friction source term must be
			 * applied as well. */
			/* can qTemp be eliminated? */
			Array2Xs qTemp = qBlock;
			auto qOld = vBlock.rowwise() * dBlock;
			auto qDif = qBlock - qOld;
			qBlock = ( dBlock < minWaterDepth || vNorm == 0 )
				.template replicate<2,1>().select(
					qTemp, // use qOld here instead? in the case where qTemp is used, they're identical, right?
					qOld + 
					( qDif - fricRepl * vBlock ) /
					( 1 + chezyRepl * ( vNormRepl + vBlock.pow(2) / vNormRepl ) )
				);

			#endif
		};


		/* there are several attempts in here, including just a cellwise loop
		 * wrapped in this blockwise call. That one is the fastest of the options
		 * here */
		// auto applyFrictionBlockwise = [&](
		// 	const auto& n, Index colBeg, Index nCols
		// ){
		// 	// auto qBlock = q().values().template bottomRows<2>().middleCols(colBeg, nCols);

		// 	// Eigen::Block<ArrayXXs, Dynamic, Dynamic, true>
		// 	// 	dBlock = d().values().middleCols(colBeg, nCols),
		// 	// 	vBlock = v().values().middleCols(colBeg, nCols);

		// 	// auto nBlock =  n.values().middleCols(colBeg, nCols);
			

		// 	// Array1Xs vNorm = vBlock.matrix().colwise().norm().array();

		// 	// /* this factor is g / C^2, with Chezy coeff. C = d^(1/6) / n */
		// 	// auto chezyFac = phys::g * nBlock * nBlock / dBlock.pow(1./3);

		// 	// /* by subtracting the old storage values from the current ones, we
		// 	//  * get all the fluxes and sources to use for point-implicit 
		// 	//  * treatment. However, first, the friction source term must be
		// 	//  * applied as well. */
		// 	// auto qOld = dBlock.replicate<2,1>() * vBlock;
		// 	// auto qDiff = qBlock - qOld;
		// 	// auto frictionSource =
		// 	// 	dt * (chezyFac * vNorm).template replicate<2,1>() * vBlock;
			
		// 	// // std::cout << type_name<decltype(frictionSource)>() << "\n";

		// 	// /* now the point-implicit treatment can be applied */
		// 	// auto qNew = qOld + (qDiff - frictionSource) / (
		// 	// 	1 + dt * (chezyFac / dBlock).template replicate<2,1>() *
		// 	// 	( vNorm.replicate<2, 1>() + vBlock.pow(2) / vNorm.replicate<2,1>() )
		// 	// );
		// 	// qBlock = (dBlock > minWaterDepth).replicate<2,1>().select(qNew, qBlock);

		// 	// // Array1Xs qNew (1,nCols);
		// 	// // for (int i=0; i<2; ++i){
		// 	// // 	auto qRow = qBlock.row(i);
		// 	// // 	auto vRow = vBlock.row(i);
		// 	// // 	auto qOld = dBlock * vRow;
		// 	// // 	auto qDif = qRow - qOld;
		// 	// // 	auto fric = dt * (chezyFac * vNorm) * vRow;
		// 	// // 	qNew = qOld + (qDif - fric) / (
		// 	// // 		1 + dt * (chezyFac / dBlock) *
		// 	// // 		( vNorm + vRow.pow(2) / vNorm )
		// 	// // 	);
		// 	// // 	qRow = (dBlock > minWaterDepth).select(qNew, qRow );
		// 	// // }

		// 	for (Index cell {colBeg}; cell<colBeg+nCols; ++cell){
		// 		if ( isDry()(0, cell) )
		// 			return;
		// 		scalar d { this->d().values()(0, cell) };
		// 		const auto& v { this->v().values().col(cell) };
		// 		scalar chezy { hms::chezy( d, n.values()(0, cell) )};
		// 		scalar fac { phys::g / (chezy * chezy) };
		// 		scalar vNorm { v.matrix().norm() };
		// 		Array2s qOld { d * v };
		// 		Array2s qDiff { q().values().col(cell).tail(2) - qOld };

		// 		/* apply friction source term */
		// 		qDiff -= dt * fac * vNorm * v;

		// 		q().values().col(cell).tail(2) = qOld + 
		// 			qDiff / ( 1 + dt * fac / d * ( vNorm + v.pow(2) / vNorm ) );
		// 	}
		// };

		if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
			static constexpr Index ng {1};
			Index
				nx = this->mesh().nCellsX(),
				ny = this->mesh().nCellsY(),
				cellsPerRow { 2*ng + nx },
				domainCellsBeg { ng*(cellsPerRow+1) };
			
			OMP_PRAGMA( parallel for )
			for ( int j = 0; j < ny; ++j ){
				std::visit( [&](const auto& n){
					applyFrictionBlockwise(n, domainCellsBeg + j*cellsPerRow, nx );
				}, n() );
			}
		} else {
			OMP_PRAGMA( parallel )
			{

			auto [colBeg, nCols] = ompSegment( this->mesh().nDomainCells() );

			std::visit( [&](const auto& n){
				applyFrictionBlockwise(n, colBeg, nCols );
			}, n() );

			}
		}

		// /* std::visit inside of the lambda */
		// auto applyFriction = [&](Index cell){
		// 	if ( isDry()(0, cell) )
		// 		return;
		// 	scalar d { this->d().values()(0, cell) };
		// 	const auto& v { this->v().values().col(cell) };
		// 	scalar chezy { hms::chezy(
		// 		d, std::visit(
		// 			[&](const auto& n){ return n.values()(0, cell); },
		// 			n()
		// 		)
		// 	)};
		// 	scalar fac { phys::g / (chezy * chezy) };
		// 	scalar vNorm { v.matrix().norm() };
		// 	Array2s qOld { d * v };
		// 	Array2s qDiff { q().values().col(cell).tail(2) - qOld };

		// 	/* apply friction source term */
		// 	qDiff -= dt * fac * vNorm * v;

		// 	q().values().col(cell).tail(2) = qOld + 
		// 		qDiff / ( 1 + dt * fac / d * ( vNorm + v.pow(2) / vNorm ) );
		// };
		
		// this->mesh().forEachDomainCellIndex( applyFriction );


		// /* std::visit outside of the lambda */
		// auto applyFriction = [&](const auto& n){
		// 	this->mesh().forEachDomainCellIndex( [&](Index cell){
		// 		if ( isDry()(0, cell) )
		// 			return;
		// 		scalar d { this->d().values()(0, cell) };
		// 		const auto& v { this->v().values().col(cell) };
		// 		scalar chezy { hms::chezy( d, n.values()(0, cell) ) };
		// 		scalar fac { phys::g / (chezy * chezy) };
		// 		scalar vNorm { v.matrix().norm() };
		// 		Array2s qOld { d * v };
		// 		Array2s qDiff { q().values().col(cell).tail(2) - qOld };

		// 		/* apply friction source term */
		// 		qDiff -= dt * fac * vNorm * v;

		// 		q().values().col(cell).tail(2) = qOld + 
		// 			qDiff / ( 1 + dt * fac / d * ( vNorm + v.pow(2) / vNorm ) );
		// 	} );
		// };
		
		// std::visit( applyFriction, n() );
	}
}

template class SWESolver<Mesh>;
template class SWESolver<UniMesh>;
template class SWESolver<RectMesh>;
template class SWESolver<StructMesh>;

} // end namespace hms
